import tensorflow as tf
from fetchData import *
import os
import sys
import cv2
import tempfile
import matplotlib.pyplot as plt

tempfile.tempdir="/home/jh/working_pros/binary_classifier/testIMG_tmp"

data_dir = "/home/jh/working_data/idCard_whole_extract"

if __name__ == "__main__":
    tf.logging.set_verbosity(tf.logging.INFO)

    idCard_data = binary_data( data_dir )
    test_data_ops  = idCard_data.testDataStream ( 128 )

    sess = tf.Session()

    saver = tf.train.import_meta_graph( "./tmp/model.ckpt.meta" )
    saver.restore( sess , tf.train.latest_checkpoint("./tmp" ) )

    graph = tf.get_default_graph()
    inputs = graph.get_tensor_by_name( "module/hub_input/images:0" )
    labels = graph.get_tensor_by_name( "label:0" )

    outputIMG = tf.map_fn( lambda img: \
            tf.image.convert_image_dtype( img , tf.uint8 ) , inputs , \
            dtype = tf.uint8 )

    train_argmax = graph.get_tensor_by_name( "train_argmax:0" )
    
    for i in range(11):
        train_images , train_labels = sess.run( test_data_ops )
        REAL_IMG, IMG, ARGMAX, LABELS = sess.run( \
                [ outputIMG , inputs,  train_argmax, labels ] , \
                feed_dict = { inputs: train_images , \
                labels : train_labels } )
        for index in range( ARGMAX.shape[0] ):
            if ARGMAX[index] != LABELS[index]:
                if LABELS[index] == 0:
                    writeFile = tempfile.mktemp( suffix=".jpg" , \
                        dir = "/home/jh/working_pros/binary_classifier/testIMG_tmp/pos")
                else:
                    writeFile = tempfile.mktemp( suffix=".jpg" , \
                        dir = "/home/jh/working_pros/binary_classifier/testIMG_tmp/neg")
                REAL_IMG[index] = cv2.cvtColor( REAL_IMG[index] , cv2.COLOR_BGR2RGB)
                cv2.imwrite( writeFile , REAL_IMG[index] )
                
                #plt.imshow( REAL_IMG[index] )
                #plt.show()
